import { generate, parser } from '@shaderfrog/glsl-parser';
import { IdentifierNode, KeywordNode, LayoutQualifierNode, Program, visit } from '@shaderfrog/glsl-parser/ast';
import { IFagmentOutLayout, IShaderLayoutInfo, IUniformLayout, IVaryingLayout, IVertexInLayout } from './IShaderLayoutInfo';
import { ShaderCodeInliner } from './inline/shaderCodeInliner';
import { isGlslVertex } from './isGlsl100';
import { preprocessor } from './preprocessor';
import { removeUnuse } from './programHandlers/removeUnuse';
import { activeUniforms } from './programHandlers/activeUniforms';
import { dFdy } from './utils/dFdy';

const head = `
#define gl_VertexID gl_VertexIndex
#define gl_InstanceID gl_InstanceIndex
`;

/**
 * 把顶点着色器单独的 Unifrom 放入该名称的 UBO 中。
 */
export const LeftOverVSUBOName = `LeftOverVSUBO`;
/**
 * 把片元着色器单独的 Unifrom 放入该名称的 UBO 中。
 */
export const LeftOverFSUBOName = `LeftOverFSUBO`;

/**
 * 升级 glsl300 版本的着色器代码为 glsl450。
 *
 * 1. 顶点片元着色器共同
 *      1. 450版本 绑定资源 只允许 UBO sampler image ，需要把单独的 unifrom 转换为 UBO 结构。
 *      2. 纹理 `samplerCube` `` 需要 纹理采样器分离。
 *          glsl300
 *          ```
 *              uniform samplerCube reflectionSampler;
 *              ...
 *              vec4 reflectionSample = textureCube(reflectionSampler, reflectionCoords);
 *          ```
 *          glsl450
 *          ```
 *              layout(set = 1, binding = 5) uniform sampler reflectionSamplerSampler;
                layout(set = 1, binding = 4) uniform textureCube reflectionSamplerTexture;
 *              ...
 *              vec4 reflectionSample = texture(samplerCube(reflectionSamplerTexture, reflectionSamplerSampler), reflectionCoords);
 *          ```
 *      3. 'sampler constructor must appear at point of use' 采样器构造函数必须在使用时出现。因此在传入采样器的函数需要使用内联方式，在预编译时进行展开来实现`采样器构造函数必须在使用时出现`。
 *
 * 2. 顶点着色器
 *      1. 为输入属性 `in` 添加 `layout` 如果没有的话。例如 `layout(location = 0) in vec3 position;` 。
 *      2. 为输出属性 `out` 添加 `layout` 。 例如 `layout(location = 0)  out vec3 vPositionW;` 。
 *      3. 为 `uniform` 新增 `layout` 。 例如 `layout(set = 0, binding = 0) uniform vec4 vPrimaryColor;` 。
 *
 * 3. 片元着色器
 *      1. 为输入属性 `in` 添加 `layout` 。 例如 `layout(location = 0)  in vec3 vPositionW;`; 注：与顶点着色器`out`的相同属性需要拥有相同的 `location` 值。
 *      2. 为输出颜色属性 `out` 添加 `layout` 如果没有的话。 例如 `layout(location = 0) out vec4 glFragColor;` 。
 *      3. 为 `uniform` 新增 `layout` 。 例如 `layout(set = 0, binding = 0) uniform vec4 vPrimaryColor;` 。注：与顶点着色器`uniform`的相同属性需要拥有相同的`set`与`binding`值。
 *
 *
 * @param code100 glsl100版本着色器代码。
 *
 * @see Babylon.js WebGPUEngine.inlineShaderCode
 */
export function glsl300ToGlsl450(code300: string, isVertex?: boolean, layoutInfo?: IShaderLayoutInfo, isDeleteUnuse = false)
{
    let newCode = code300;

    if (isVertex === undefined)
    {
        isVertex = isGlslVertex(code300);
    }

    //
    layoutInfo = layoutInfo || {
        uniform: { defaultSet: 0, set: [], map: {}, structType: {} },
        vertexIn: { locations: [], vertexIns: {} },
        varying: { locations: [], varyings: {} },
        fagmentOut: { locations: [], fagmentOuts: {} },
        LeftOverVSUBOName,
        LeftOverFSUBOName,
        activeUniforms: [],
    };

    // 不允许结构体中包含纹理采样
    checkTextureInStruct(code300);

    newCode = dFdy(newCode);

    // 处理纹理，分离为采样器与纹理。
    newCode = textureProcessor(newCode);

    // 预处理宏定义
    newCode = preprocessor(head + newCode);

    // 解析语法树
    let program = parser.parse(newCode);

    // 移除 多余定义的内置函数
    removeBuiltinFunction(program);

    // 移除未被使用的
    if (isDeleteUnuse)
    {
        program = removeUnuse(program);
    }

    //
    activeUniforms(program, layoutInfo);

    // 把单独的 unifrom 转换为 UBO
    uniformBlockProcessor(program, isVertex);

    // 绑定资源， 为 UBO sample texture 新增 `layout`
    layoutUniformProcessor(program, layoutInfo);

    //
    if (isVertex)
    {
        // 为 顶点输入属性 `in` 新增 `layout`
        vertextInProcessor(program, layoutInfo.vertexIn);

        // 为 顶点输出变量 `out` 新增 `layout`
        vertextOutProcessor(program, layoutInfo.varying);
    }
    else
    {
        // 为 片元输入变量 `in` 新增 `layout`
        fragmentInProcessor(program, layoutInfo.varying);

        // 为 片元输输出颜色 `out` 新增 `layout`
        fragmentOutProcessor(program, layoutInfo.fagmentOut);
    }

    // 从语法树还原为代码。
    newCode = generate(program);

    // 格式化代码
    newCode = preprocessor(newCode);

    const code450 = `#version 450\n${newCode}`;

    return { code: code450, layoutInfo };
}

/**
 * 不允许结构体中包含纹理采样
 *
 * @param code 着色器代码。
 */
function checkTextureInStruct(code: string)
{
    const structReg = /struct\s[^{]+{([^}]+)}/g;
    let result = structReg.exec(code);
    while (result)
    {
        const [str, body] = result;

        const samplerReg = /\b(usampler2D|sampler2D|sampler2DArray|sampler2DShadow|sampler2DArrayShadow|samplerCube|samplerCubeArray|sampler3D)\b/g;
        if (body.match(samplerReg))
        {
            console.error(`由于UBO不支持采样器，因此着色器结构体中包含采样器时无法升级！ \n${str}`);
        }

        result = structReg.exec(code);
    }
}

function vertextInProcessor(program: Program, attributeLayout: IVertexInLayout)
{
    const addLayouts: Function[] = [];

    program.program.forEach((element) =>
    {
        if (element.type !== 'declaration_statement') return;

        if (element.declaration.type !== 'declarator_list') return;

        const qualifiers = element.declaration.specified_type.qualifiers;

        const key0 = (qualifiers?.[0] as KeywordNode)?.token;
        const key1 = (qualifiers?.[1] as KeywordNode)?.token;

        if (!(key1 === 'in' || key0 === 'in')) return;

        const identifier = element.declaration.declarations[0].identifier.identifier;

        if (qualifiers?.[0].type === 'layout_qualifier')
        {
            const location = Number((qualifiers[0].qualifiers[0].expression as KeywordNode).token);
            // 验证 location 是否一致
            console.assert(!isNaN(location) && (attributeLayout.vertexIns[identifier] === undefined || attributeLayout.vertexIns[identifier] === location));
            //
            attributeLayout.locations[location] = identifier;
            attributeLayout.vertexIns[identifier] = location;
        }
        else
        {
            let length = 1;

            // 处理数组，如 `in mat4 v_blurTexCoords[2];`
            if (element.declaration.declarations[0].quantifier?.[0].type === 'array_specifier')
            {
                const lengthToken = (element.declaration.declarations[0].quantifier[0].expression as KeywordNode).token;
                console.assert(lengthToken !== undefined);

                length = Number(lengthToken);
            }
            const typeName = (element.declaration.specified_type.specifier.specifier as KeywordNode)?.token;
            if (typeName === 'mat4')
            {
                length = length * 4;
            }

            addLayouts.push(() =>
            {
                const location = getAttributeLayout(attributeLayout, identifier, length);

                const node = createLayoutLocationNode(location);

                qualifiers.unshift(node);
            });
        }
    });

    // 执行添加 layout
    addLayouts.forEach((f) => f());
}

function vertextOutProcessor(program: Program, varyingLayout: IVaryingLayout)
{
    program.program.forEach((element) =>
    {
        if (element.type !== 'declaration_statement') return;

        if (element.declaration.type !== 'declarator_list') return;

        const key0 = (element.declaration.specified_type.qualifiers?.[0] as KeywordNode)?.token;
        const key1 = (element.declaration.specified_type.qualifiers?.[1] as KeywordNode)?.token;

        // 支持带关键字 flat|centroid 等关键字 的顶点输出属性。
        if (!(key1 === 'out' || key0 === 'out')) return;

        const identifier = element.declaration.declarations[0].identifier.identifier;
        //
        let length = 1;
        // 处理数组，如 `out vec2 v_blurTexCoords[14];`
        if (element.declaration.declarations[0].quantifier?.[0].type === 'array_specifier')
        {
            const lengthToken = (element.declaration.declarations[0].quantifier[0].expression as KeywordNode).token;
            console.assert(lengthToken !== undefined);

            length = Number(lengthToken);
        }

        const location = getVaryingLayout(varyingLayout, identifier, length);

        const node = createLayoutLocationNode(location);

        element.declaration.specified_type.qualifiers.unshift(node);
    });
}

/**
 * 为 片元输入变量 `in` 新增 `layout` 。
 *
 * @param program 语法树。
 */
function fragmentInProcessor(program: Program, varyingLayout: IVaryingLayout)
{
    program.program.forEach((element) =>
    {
        if (element.type !== 'declaration_statement') return;

        if (element.declaration.type !== 'declarator_list') return;

        if (!element.declaration.specified_type.qualifiers) return;

        if (element.declaration.specified_type.qualifiers[0].type !== 'keyword') return;

        const qualifiers = element.declaration.specified_type.qualifiers;
        const key0 = (qualifiers?.[0] as KeywordNode)?.token;
        const key1 = (qualifiers?.[1] as KeywordNode)?.token;

        // 支持带关键字 flat|centroid 等关键字 的顶点输出属性。
        if (!(key1 === 'in' || key0 === 'in')) return;

        const identifier = element.declaration.declarations[0].identifier.identifier;

        let length = 1;
        // 处理数组，如 `in vec2 v_blurTexCoords[14];`
        if (element.declaration.declarations[0].quantifier?.[0].type === 'array_specifier')
        {
            const lengthToken = (element.declaration.declarations[0].quantifier[0].expression as KeywordNode).token;
            console.assert(lengthToken !== undefined);

            length = Number(lengthToken);
        }

        const location = getVaryingLayout(varyingLayout, identifier, length);
        const node = createLayoutLocationNode(location);

        qualifiers.unshift(node);
    });
}

/**
 * 为 片元输输出颜色 `out` 新增 `layout` 。
 *
 * @param program 语法树。
 */
function fragmentOutProcessor(program: Program, fagmentOutLayout: IFagmentOutLayout)
{
    const addLayouts: Function[] = [];

    program.program.forEach((element) =>
    {
        if (element.type !== 'declaration_statement') return;

        if (element.declaration.type !== 'declarator_list') return;

        if (!element.declaration.specified_type.qualifiers) return;

        const qualifiers = element.declaration.specified_type.qualifiers;
        const key0 = (qualifiers?.[0] as KeywordNode)?.token;
        const key1 = (qualifiers?.[1] as KeywordNode)?.token;

        if (!(key1 === 'out' || key0 === 'out')) return;

        const identifier = element.declaration.declarations[0].identifier.identifier;

        if (qualifiers[0].type === 'layout_qualifier')
        {
            const location = Number((qualifiers[0].qualifiers[0].expression as KeywordNode).token);
            // 验证 location 是否一致
            console.assert(!isNaN(location) && (fagmentOutLayout.fagmentOuts[identifier] === undefined || fagmentOutLayout.fagmentOuts[identifier] === location));
            //
            fagmentOutLayout.locations[location] = identifier;
            fagmentOutLayout.fagmentOuts[identifier] = location;
        }
        else
        {
            addLayouts.push(() =>
            {
                const location = getFagmentOutLayout(fagmentOutLayout, identifier);
                // 为 out 新增 layout
                const node = createLayoutLocationNode(location);
                qualifiers.unshift(node);
            });
        }
    });

    //
    addLayouts.forEach((f) => f());
}

/**
 * 为 `unifrom` 新增 `layout` , 直接在语法树中新增 `layout` 相关节点。
 *
 * @param program 通过 `glsl-parser` 解析的程序语法树。
 */
function layoutUniformProcessor(program: Program, options: IShaderLayoutInfo)
{
    program.program.forEach((element) =>
    {
        if (element.type !== 'declaration_statement') return;

        // 处理 UBO
        if (element.declaration.type === 'interface_declarator')
        {
            if ((element.declaration.qualifiers[0] as KeywordNode)?.token !== 'uniform') return;
            //
            let identifier = element.declaration.identifier?.identifier?.identifier;
            const structType = element.declaration.interface_type?.identifier;

            if (!identifier)
            {
                identifier = structType;
            }

            options.uniform.structType[structType] = identifier;
            //
            const { set, binding } = getSetBinding(options.uniform, identifier);

            const layoutQualifierNode = createLayoutNode(set, binding);

            element.declaration.qualifiers.unshift(layoutQualifierNode);
        }
        else if (element.declaration.type === 'declarator_list')
        {
            if ((element.declaration.specified_type.qualifiers?.[0] as KeywordNode)?.token !== 'uniform') return;

            // 处理 sampler texture
            const identifier = element.declaration.declarations[0].identifier.identifier;

            const { set, binding } = getSetBinding(options.uniform, identifier);

            const layoutQualifierNode = createLayoutNode(set, binding);
            element.declaration.specified_type.qualifiers.unshift(layoutQualifierNode);
        }
    });
}

const builtinFuncNames = ['transpose', 'inverse'];
/**
 * 移除 多余定义的内置函数。
 *
 * @param program 语法树。
 */
function removeBuiltinFunction(program: Program)
{
    for (let i = program.program.length - 1; i >= 0; i--)
    {
        // 遍历程序查找 unifrom
        const element = program.program[i];
        if (element.type === 'function')
        {
            const funcName = element.prototype.header.name.identifier;
            if (builtinFuncNames.indexOf(funcName) !== -1)
            {
                program.program.splice(i, 1);
            }
        }
    }
    const functionScopeIndex = program.scopes[0].functions;
    const funcNames = Object.keys(functionScopeIndex);
    for (let i = 0; i < funcNames.length; i++)
    {
        const funcName = funcNames[i];
        if (builtinFuncNames.indexOf(funcName) !== -1)
        {
            delete functionScopeIndex[funcName];
            program.scopes = program.scopes.filter((v) => v.name !== funcName);
        }
    }
}

/**
 * 把单独的 unifrom 转换为 UBO。
 *
 * @param program 语法树。
 */
function uniformBlockProcessor(program: Program, isVertex: boolean)
{
    let unifroms = '';
    let first = -1;

    for (let i = program.program.length - 1; i >= 0; i--)
    {
        // 遍历程序查找 unifrom
        const element = program.program[i];
        if (element.type !== 'declaration_statement') continue;
        if (element.declaration.type !== 'declarator_list') continue;
        if ((element.declaration.specified_type.qualifiers?.[0] as KeywordNode)?.token !== 'uniform') continue;
        const type = (element.declaration.specified_type.specifier.specifier as IdentifierNode).identifier;
        // 跳过采样器与纹理
        if (type && (type.indexOf('sampler') !== -1 || type.indexOf('texture') !== -1)) continue;

        // 删除关键字 uniform
        element.declaration.specified_type.qualifiers.shift();

        const code = generate(element);

        unifroms = code + unifroms;

        program.program.splice(i, 1);
        first = i;
    }

    // 把 uniform 组装为 UBO
    if (first > -1)
    {
        const unifromsProgram = parser.parse(`
uniform ${isVertex ? LeftOverVSUBOName : LeftOverFSUBOName}
{
    ${unifroms}
};`);

        program.program.splice(first, 0, unifromsProgram.program[0]);
    }
}

function getAttributeLayout(attributeLayout: IVertexInLayout, identifier: string, length: number)
{
    let location = attributeLayout.vertexIns[identifier];

    if (location === undefined)
    {
        // 获取第一个空位置
        location = findFirstUndefinedIndex(attributeLayout.locations, length);
        //
        for (let i = 0; i < length; i++)
        {
            attributeLayout.locations[location + i] = identifier;
        }
        //
        attributeLayout.vertexIns[identifier] = location;
    }

    return location;
}

function getVaryingLayout(varyingLayout: IVaryingLayout, identifier: string, length: number)
{
    let location = varyingLayout.varyings[identifier];

    if (location === undefined)
    {
        // 获取第一个空位置
        location = findFirstUndefinedIndex(varyingLayout.locations, length);
        //
        for (let i = 0; i < length; i++)
        {
            varyingLayout.locations[location + i] = identifier;
        }
        varyingLayout.varyings[identifier] = location;
    }

    return location;
}

function getFagmentOutLayout(fagmentOutLayout: IFagmentOutLayout, identifier: string)
{
    let location = fagmentOutLayout.fagmentOuts[identifier];

    if (location === undefined)
    {
        // 获取第一个空位置
        location = findFirstUndefinedIndex(fagmentOutLayout.locations);
        //
        fagmentOutLayout.locations[location] = identifier;
        fagmentOutLayout.fagmentOuts[identifier] = location;
    }

    return location;
}

/**
 * 获取数组中第一个连续 length 个 undefined 位置。
 *
 * @param array 数组
 * @param length 数量。
 * @returns 出现的位置。
 */
function findFirstUndefinedIndex(array: any[], length = 1)
{
    let result = 0;
    let undefinedNum = 0;
    for (let i = 0; i < 1000; i++)
    {
        if (array[i] === undefined)
        {
            undefinedNum++;
        }
        else
        {
            undefinedNum = 0;
        }
        if (undefinedNum >= length)
        {
            result = i - length + 1;
            break;
        }
    }

    return result;
}

function getSetBinding(uniformSetBinding: IUniformLayout, identifier: string)
{
    let set = 0;
    let binding = 0;

    const setBinding = uniformSetBinding.map[identifier];

    if (!setBinding)
    {
        //
        set = uniformSetBinding.defaultSet;
        uniformSetBinding.set[set] = uniformSetBinding.set[set] || [];

        binding = uniformSetBinding.set[set].length;
        //
        uniformSetBinding.set[set][binding] = identifier;
        uniformSetBinding.map[identifier] = { set, binding };
    }
    else
    {
        set = setBinding.set;
        binding = setBinding.binding;
    }

    return { set, binding };
}

/**
 * 升级 glsl300 中纹理使用方式。
 *
 * @param code 着色器代码。
 * @returns 升级后的着色器代码。
 *
 * 注：'sampler constructor must appear at point of use' 采样器构造函数必须在使用时出现。因此在传入采样器的函数需要使用内联方式，在预编译时进行展开来实现`采样器构造函数必须在使用时出现`。
 *
 * @see Babylon.js WebGPUShaderProcessorGLSL.uniformProcessor
 * @see Babylon.js WebGPUEngine.inlineShaderCode
 */
function textureProcessor(code: string)
{
    let newCode = code;

    // uniform 正则
    const uniformRegex = /\s*(uniform\s+((?:highp)?|(?:lowp)?|(?:mediump)?)\s*(\S+)\s+(\S+)\s*;)/g;

    // 是否存在函数名称为 'texture' 纹理。
    let hasTextureName = false;
    let replaceTextureDefineStr: string;

    //
    let match = uniformRegex.exec(newCode);
    while (match)
    {
        const str = match[1];
        const uniformPrecision = match[2];
        let uniformType = match[3];
        const name = match[4];

        if (uniformType.indexOf('sampler') === 0 || uniformType.indexOf('sampler') === 1)
        {
            const componentType = uniformType.charAt(0) === 'u' ? 'u' : uniformType.charAt(0) === 'i' ? 'i' : '';
            if (componentType)
            {
                uniformType = uniformType.substring(1);
            }
            const samplerType = _SamplerTypeByWebGLSamplerType[uniformType] ?? 'sampler';
            const samplerName = name + AutoSamplerSuffix;
            const textureType = _TextureTypeByWebGLSamplerType[uniformType];
            const samplerFunction = _SamplerFunctionByWebGLSamplerType[uniformType];

            const replaceDefineStr = `${componentType}${samplerFunction}(${name}Texture, ${samplerName})`;

            const replaceStr = `uniform ${samplerType} ${samplerName};
                        uniform ${componentType}${textureType} ${name}Texture;
                        #define ${name} ${replaceDefineStr}`;

            newCode = newCode.substring(0, uniformRegex.lastIndex - str.length)
                + replaceStr
                + newCode.substring(uniformRegex.lastIndex);

            uniformRegex.lastIndex += replaceStr.length - str.length;

            // 处理纹理名称与函数名称同名 'texture' 情况。
            if (name === 'texture')
            {
                hasTextureName = true;
                replaceTextureDefineStr = replaceDefineStr;
            }
        }

        //
        match = uniformRegex.exec(newCode);
    }

    // 标记所有使用采样器的函数为内敛函数，解决`采样器构造函数必须在使用时出现`问题。
    const functionReg = /\S+\s+[\w+]+\(([^)]+)\)\s*\{/g;
    const samplerReg = /sampler/g;
    const defineInlineStr = `#define inline\n`;
    match = functionReg.exec(newCode);
    while (match)
    {
        // 判断参数中是否存在采样器
        const [_, params] = match;

        if (params.match(samplerReg))
        {
            // 在函数前新增内联标记
            newCode = newCode.substring(0, functionReg.lastIndex - _.length) + defineInlineStr + newCode.substring(functionReg.lastIndex - _.length);
            functionReg.lastIndex += defineInlineStr.length;
        }

        match = functionReg.exec(newCode);
    }

    // 处理函数内联
    const sci = new ShaderCodeInliner(newCode);
    sci.debug = false;
    sci.processCode();
    newCode = sci.code;

    // 处理 存在 纹理名称为 'texture' 情况。
    if (hasTextureName)
    {
        newCode = preprocessor(newCode);

        let index = newCode.indexOf(replaceTextureDefineStr);
        while (index !== -1)
        {
            for (let i = index + replaceTextureDefineStr.length; i < newCode.length; i++)
            {
                const char = newCode.charAt(i);
                if (char.match(/\s/)) continue;
                if (char === '(')
                {
                    newCode = `${newCode.substring(0, index)}texture${newCode.substring(i)}`;
                }
                break;
            }
            index = newCode.indexOf(replaceTextureDefineStr, index + replaceTextureDefineStr.length);
        }
    }

    return newCode;
}

const _SamplerTypeByWebGLSamplerType = {
    sampler2DShadow: 'samplerShadow',
    sampler2DArrayShadow: 'samplerShadow',
};
const AutoSamplerSuffix = 'Sampler';
const _TextureTypeByWebGLSamplerType = {
    sampler2D: 'texture2D',
    sampler2DArray: 'texture2DArray',
    sampler2DShadow: 'texture2D',
    sampler2DArrayShadow: 'texture2DArray',
    samplerCube: 'textureCube',
    samplerCubeArray: 'textureCubeArray',
    sampler3D: 'texture3D',
};
const _SamplerFunctionByWebGLSamplerType = {
    sampler2D: 'sampler2D',
    sampler2DArray: 'sampler2DArray',
    sampler2DShadow: 'sampler2DShadow',
    sampler2DArrayShadow: 'sampler2DArrayShadow',
    samplerCube: 'samplerCube',
    sampler3D: 'sampler3D',
};

function createLayoutNode(set: number, binding: number)
{
    const code = `layout(set = ${set}, binding = ${binding});`;
    const program = parser.parse(code);

    let result: LayoutQualifierNode;

    visit(program, {
        layout_qualifier: {
            enter: (path) =>
            {
                result = path.node;
                const code = generate(path.node);

                code;
            }
        },
    });

    return result;
}

function createLayoutLocationNode(location: number)
{
    const code = `layout(location = ${location});`;
    const program = parser.parse(code);

    let result: LayoutQualifierNode;

    visit(program, {
        layout_qualifier: {
            enter: (path) =>
            {
                result = path.node;
                const code = generate(path.node);

                code;
            }
        },
    });

    return result;
}
